import uuid
import psycopg2
from psycopg2 import OperationalError, ProgrammingError
from psycopg2.extras import DictCursor  # 字典游标，保持返回格式兼容

class HistoryMemory:
    """基于 PostgreSQL 的对话历史管理类（适配原生 UUID 类型与触发器）"""

    def __init__(self, db_config, max_sessions=10, max_turns_per_session=20):
        self.max_sessions = max_sessions
        self.max_turns_per_session = max_turns_per_session
        self.db_config = db_config  # 配置格式：{host, user, password, database}
        self.current_user_id = None
        self.current_session_id = None
        self._create_necessary_tables()  # 初始化表与触发器

    def _get_db_connection(self):
        """获取 PostgreSQL 连接（替换 PyMySQL 连接逻辑）"""
        try:
            conn = psycopg2.connect(
                host=self.db_config['host'],
                user=self.db_config['postgres'],
                password=self.db_config['ht2025'],
                dbname=self.db_config['database'],  # PostgreSQL 用 dbname，非 database
                cursor_factory=DictCursor  # 字典游标，返回 {字段名: 值} 格式
            )
            return conn
        except OperationalError as e:
            print(f"❌ 数据库连接失败: {str(e)}")
        return None

    def _create_necessary_tables(self):
        """创建表与触发器（PostgreSQL 语法适配）"""
        conn = self._get_db_connection()
        if not conn:
            return
        cursor = conn.cursor()
        try:
            # 1. 创建会话表（id 用原生 UUID 类型，添加触发器实现 updated_at 自动更新）
            create_session_sql = """
            CREATE TABLE IF NOT EXISTS chat_sessions (
                id UUID NOT NULL COMMENT '会话ID（原生UUID）',
                user_id INT NOT NULL COMMENT '关联用户ID',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,  # 初始值为创建时间
                PRIMARY KEY (id),
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE,
                INDEX idx_user_id (user_id)  # 用户ID索引，加速查询
            );
            """
            cursor.execute(create_session_sql)

            # 2. 创建对话轮次表（session_id 与会话表 UUID 类型对齐）
            create_turn_sql = """
            CREATE TABLE IF NOT EXISTS chat_turns (
                id UUID NOT NULL COMMENT '轮次ID（原生UUID）',
                session_id UUID NOT NULL COMMENT '关联会话ID（与chat_sessions.id对齐）',
                user_input TEXT NOT NULL,
                agent_response TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (id),
                FOREIGN KEY (session_id) REFERENCES chat_sessions(id) ON DELETE CASCADE,
                INDEX idx_session_id (session_id)  # 会话ID索引，加速查询
            );
            """
            cursor.execute(create_turn_sql)

            # 3. 创建触发器函数：更新时自动刷新 updated_at
            create_trigger_func_sql = """
            CREATE OR REPLACE FUNCTION update_updated_at_column()
            RETURNS TRIGGER AS $$
            BEGIN
                NEW.updated_at = CURRENT_TIMESTAMP;  # 更新时重置为当前时间
                RETURN NEW;
            END;
            $$ LANGUAGE plpgsql;
            """
            cursor.execute(create_trigger_func_sql)

            # 4. 为 chat_sessions 表添加触发器（避免重复创建）
            create_trigger_sql = """
            CREATE TRIGGER IF NOT EXISTS trigger_chat_sessions_update
            BEFORE UPDATE ON chat_sessions
            FOR EACH ROW
            EXECUTE FUNCTION update_updated_at_column();
            """
            cursor.execute(create_trigger_sql)

            conn.commit()
            print("✅ 表与触发器初始化成功（PostgreSQL）")
        except ProgrammingError as e:
            conn.rollback()
            print(f"❌ 创建表/触发器失败: {str(e)}（表已存在可忽略）")
        finally:
            cursor.close()
            conn.close()

    def _get_or_create_session(self, session_id=None, user_id=None):
        """内部创建/获取会话（适配 PostgreSQL UUID 类型与字典游标）"""
        user_id = user_id or self.current_user_id
        try:
            user_id = int(user_id)
        except (ValueError, TypeError):
            raise ValueError(f"无效用户ID: {user_id}")
        if not user_id:
            raise ValueError("用户ID不能为空")

        conn = self._get_db_connection()
        if not conn:
            raise Exception("数据库连接失败")

        cursor = conn.cursor()
        try:
            # 生成 UUID 会话ID（若未传入）
            if not session_id:
                session_id = str(uuid.uuid4())  # PostgreSQL 可直接接收 UUID 字符串

            # 检查会话是否已存在
            check_sql = "SELECT id FROM chat_sessions WHERE user_id = %s AND id = %s;"
            cursor.execute(check_sql, (user_id, session_id))
            existing_session = cursor.fetchone()

            if not existing_session:
                # 插入新会话（id 为 UUID 字符串，PostgreSQL 自动解析为 UUID 类型）
                insert_sql = "INSERT INTO chat_sessions (id, user_id) VALUES (%s, %s);"
                cursor.execute(insert_sql, (session_id, user_id))
                print(f"📥 新建会话: user_id={user_id}, session_id={session_id}")

                # 限制会话数量（超过 max_sessions 时删除最早会话）
                list_sql = "SELECT id FROM chat_sessions WHERE user_id = %s ORDER BY created_at ASC;"
                cursor.execute(list_sql, (user_id,))
                all_sessions = cursor.fetchall()  # 字典列表：[{id: ...}, ...]

                if len(all_sessions) > self.max_sessions:
                    oldest_session_id = all_sessions[0]['id']  # 字典游标直接取 id 字段
                    delete_sql = "DELETE FROM chat_sessions WHERE id = %s;"
                    cursor.execute(delete_sql, (oldest_session_id,))
                    print(f"🗑️ 删除最早会话: {oldest_session_id}")

            self.current_user_id = user_id
            self.current_session_id = session_id
            conn.commit()
            return session_id
        except ProgrammingError as e:
            conn.rollback()
            raise Exception(f"操作会话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()

    def start_new_session(self, user_id=None):
        """创建新会话（接口逻辑不变）"""
        user_id = user_id or self.current_user_id
        if not user_id:
            raise ValueError("需指定用户ID")
        return self._get_or_create_session(user_id=user_id)

    def add_turn(self, user_input: str, agent_response: str, session_id=None, user_id=None):
        """添加对话轮次（适配 PostgreSQL UUID 外键与批量删除）"""
        user_id = user_id or self.current_user_id
        session_id = session_id or self.current_session_id
        if not (user_id and session_id and user_input.strip() and agent_response.strip()):
            raise ValueError("参数不能为空")
        try:
            user_id = int(user_id)
        except (ValueError, TypeError):
            raise ValueError(f"无效用户ID: {user_id}")

        conn = self._get_db_connection()
        if not conn:
            raise Exception("数据库连接失败")

        cursor = conn.cursor()
        try:
            # 检查会话归属（避免非法会话写入）
            check_sql = "SELECT id FROM chat_sessions WHERE user_id = %s AND id = %s;"
            cursor.execute(check_sql, (user_id, session_id))
            if not cursor.fetchone():
                raise Exception(f"会话不存在或不属于该用户: session_id={session_id}")

            # 限制单会话轮次数量（超过时删除最早轮次）
            list_sql = "SELECT id FROM chat_turns WHERE session_id = %s ORDER BY created_at ASC;"
            cursor.execute(list_sql, (session_id,))
            all_turns = cursor.fetchall()  # 字典列表：[{id: ...}, ...]

            if len(all_turns) >= self.max_turns_per_session:
                # 计算需删除的多余轮次（保留最后 max_turns_per_session 轮）
                excess_turns = all_turns[:-self.max_turns_per_session]
                excess_ids = [turn['id'] for turn in excess_turns]  # 提取 UUID 列表

                # 批量删除（PostgreSQL 支持 IN 子句批量参数）
                placeholders = ", ".join(["%s"] * len(excess_ids))
                delete_sql = f"DELETE FROM chat_turns WHERE id IN ({placeholders});"
                cursor.execute(delete_sql, excess_ids)
                print(f"🗑️ 删除{len(excess_ids)}轮旧对话（会话: {session_id}）")

            # 插入新对话轮次（生成 UUID 轮次ID）
            turn_id = str(uuid.uuid4())
            insert_sql = """
            INSERT INTO chat_turns (id, session_id, user_input, agent_response)
            VALUES (%s, %s, %s, %s);
            """
            cursor.execute(insert_sql, (turn_id, session_id, user_input, agent_response))
            conn.commit()
            print(f"✅ 保存对话: turn_id={turn_id[:8]}...（会话: {session_id}）")
        except ProgrammingError as e:
            conn.rollback()
            raise Exception(f"保存对话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()

    def get_history(self, session_id=None, user_id=None) -> list:
        """获取对话历史（适配 PostgreSQL TO_CHAR 日期格式化）"""
        user_id = user_id or self.current_user_id
        session_id = session_id or self.current_session_id
        try:
            user_id = int(user_id) if user_id else None
        except (ValueError, TypeError):
            print(f"⚠️ 无效用户ID: {user_id}")
            return []
        if not (user_id and session_id):
            print(f"⚠️ 参数为空: user_id={user_id}, session_id={session_id}")
            return []

        conn = self._get_db_connection()
        if not conn:
            return []

        cursor = conn.cursor()
        try:
            # PostgreSQL 用 TO_CHAR 格式化日期（替代 MySQL DATE_FORMAT）
            get_history_sql = """
            SELECT 
                id AS turn_id, 
                user_input AS "user",  # "user" 是关键字，需加引号
                agent_response AS agent, 
                TO_CHAR(created_at, 'YYYY-MM-DD HH24:MI:SS') AS timestamp  # 24小时制格式
            FROM chat_turns 
            WHERE session_id = %s 
            ORDER BY created_at ASC;
            """
            cursor.execute(get_history_sql, (session_id,))
            history = cursor.fetchall()  # 字典列表：[{turn_id: ..., user: ...}, ...]
            print(f"📥 获取到{len(history)}轮对话（会话: {session_id}）")

            # 格式化返回结果（与原逻辑一致）
            return [
                {
                    "id": turn["turn_id"],
                    "timestamp": turn["timestamp"],
                    "user": turn["user"],
                    "agent": turn["agent"]
                }
                for turn in history
            ]
        except ProgrammingError as e:
            print(f"❌ 获取历史失败: {str(e)} | SQL: {get_history_sql} | 参数: ({session_id},)")
            return []
        finally:
            cursor.close()
            conn.close()

    def get_sessions_list(self, user_id=None) -> list:
        """获取会话列表（适配 PostgreSQL TO_CHAR 与字典游标）"""
        user_id = user_id or self.current_user_id
        try:
            user_id = int(user_id)
        except (ValueError, TypeError):
            print(f"⚠️ 无效用户ID: {user_id}")
            return []
        if not user_id:
            return []

        conn = self._get_db_connection()
        if not conn:
            return []

        cursor = conn.cursor()
        try:
            # 获取用户所有会话（按创建时间倒序）
            get_sessions_sql = """
            SELECT 
                id AS session_id, 
                TO_CHAR(created_at, 'YYYY-MM-DD HH24:MI:SS') AS created_at
            FROM chat_sessions 
            WHERE user_id = %s 
            ORDER BY created_at DESC;
            """
            cursor.execute(get_sessions_sql, (user_id,))
            sessions = cursor.fetchall()  # 字典列表：[{session_id: ..., created_at: ...}, ...]
            print(f"📥 获取到{len(sessions)}个会话（用户: {user_id}）")

            # 为每个会话添加标题（取第一条对话的前30字符）
            session_list = []
            for session in sessions:
                session_id = session["session_id"]
                created_at = session["created_at"]

                # 查询会话第一条对话（作为标题）
                get_title_sql = """
                SELECT user_input 
                FROM chat_turns 
                WHERE session_id = %s 
                ORDER BY created_at ASC 
                LIMIT 1;
                """
                cursor.execute(get_title_sql, (session_id,))
                first_turn = cursor.fetchone()

                # 生成标题（无对话则显示"新会话"）
                if first_turn and first_turn["user_input"]:
                    title = first_turn["user_input"][:30] + "..." if len(first_turn["user_input"]) > 30 else first_turn["user_input"]
                else:
                    title = "新会话"

                session_list.append({
                    "id": session_id,
                    "title": title,
                    "timestamp": created_at
                })
            return session_list
        except ProgrammingError as e:
            print(f"❌ 获取会话列表失败: {str(e)} | SQL: {get_sessions_sql} | 参数: ({user_id},)")
            return []
        finally:
            cursor.close()
            conn.close()

    # ------------------------------ 以下方法逻辑不变，仅适配异常类型 ------------------------------
    def get_context(self, session_id=None, user_id=None, last_n=5) -> str:
        history = self.get_history(session_id, user_id)
        recent = history[-last_n:] if len(history) >= last_n else history
        parts = []
        for turn in recent:
            parts.append(f"用户: {turn['user']}")
            parts.append(f"助手: {turn['agent']}")
        return "\n".join(parts)

    def clear_session(self, session_id=None, user_id=None):
        user_id = user_id or self.current_user_id
        session_id = session_id or self.current_session_id
        try:
            user_id = int(user_id) if user_id else None
        except (ValueError, TypeError):
            return
        if not (user_id and session_id):
            return

        conn = self._get_db_connection()
        if not conn:
            return

        cursor = conn.cursor()
        try:
            cursor.execute("DELETE FROM chat_turns WHERE session_id = %s;", (session_id,))
            conn.commit()
            print(f"🧹 清空会话{session_id}：{cursor.rowcount}轮对话")
        except ProgrammingError as e:
            conn.rollback()
            print(f"❌ 清空会话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()

    def delete_session(self, session_id, user_id=None):
        user_id = user_id or self.current_user_id
        try:
            user_id = int(user_id) if user_id else None
        except (ValueError, TypeError):
            return
        if not (user_id and session_id):
            return

        conn = self._get_db_connection()
        if not conn:
            return

        cursor = conn.cursor()
        try:
            # 先检查会话归属
            cursor.execute("SELECT id FROM chat_sessions WHERE user_id = %s AND id = %s;", (user_id, session_id))
            if not cursor.fetchone():
                print(f"⚠️ 会话不存在或不属于用户: {session_id}")
                return

            # 删除会话（外键级联删除对应对话轮次）
            cursor.execute("DELETE FROM chat_sessions WHERE user_id = %s AND id = %s;", (user_id, session_id))
            conn.commit()
            print(f"🗑️ 删除会话{session_id}（用户: {user_id}）")
            if self.current_session_id == session_id:
                self.current_session_id = None
        except ProgrammingError as e:
            conn.rollback()
            print(f"❌ 删除会话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()

    def clear_all_sessions(self, user_id=None):
        user_id = user_id or self.current_user_id
        try:
            user_id = int(user_id) if user_id else None
        except (ValueError, TypeError):
            return
        if not user_id:
            return

        conn = self._get_db_connection()
        if not conn:
            return

        cursor = conn.cursor()
        try:
            # 级联删除所有会话及对应对话
            cursor.execute("DELETE FROM chat_sessions WHERE user_id = %s;", (user_id,))
            conn.commit()
            print(f"🗑️ 删除用户{user_id}的{cursor.rowcount}个会话")
            self.current_session_id = None
        except ProgrammingError as e:
            conn.rollback()
            print(f"❌ 清空所有会话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()

    def set_current_session(self, session_id, user_id=None):
        user_id = user_id or self.current_user_id
        try:
            user_id = int(user_id) if user_id else None
        except (ValueError, TypeError):
            return
        if not (user_id and session_id):
            return

        conn = self._get_db_connection()
        if not conn:
            return

        cursor = conn.cursor()
        try:
            # 验证会话归属
            cursor.execute("SELECT id FROM chat_sessions WHERE user_id = %s AND id = %s;", (user_id, session_id))
            if cursor.fetchone():
                self.current_user_id = user_id
                self.current_session_id = session_id
                print(f"🔄 切换到会话: {session_id}（用户: {user_id}）")
            else:
                print(f"⚠️ 会话不存在或不属于用户: {session_id}")
        except ProgrammingError as e:
            print(f"❌ 切换会话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()